// SPDX-License-Identifier: GPL-2.0-only
/*
 * Example code for CSV RTMR & ATTESTATION.
 *
 * Copyright (C) 2025 Hygon Info Technologies Ltd.
 *
 * Author: Liyang Han <hanliyang@hygon.cn>
 */

#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <openssl/evp.h>

#include "csv-guest-module-tool.h"
#include "rtmr.h"

static int do_rtmr_status(void)
{
	int fd;
	int ret;
	struct csv_rtmr_req req = { 0, };
	struct csv_guest_user_rtmr_status status = { 0, };

	// access csv-guest
	fd = open_csv_guest();
	if (fd < 0)
		return -1;

	// issue RTMR_EXTEND request
	req.buf = (__u64)&status;
	req.len = (__u64)sizeof(status);
	req.subcmd_id = CSV_GUEST_USER_RTMR_STATUS;
	ret = ioctl(fd, CSV_CMD_RTMR, &req);

	// check return code
	if (ret == -ENODEV || ret == -ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		goto err;
	}

	if (errno == ENODEV || errno == ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		ret = -errno;
		goto err;
	}

	if (req.fw_error_code) {
		PR_ERR("fw_err: 0x%x\n", req.fw_error_code);
		ret = -EIO;
	}

	PR_INFO("RTMR_STATUS:\n");
	PR_INFO("  version: 0x%x\n", status.version);
	PR_INFO("  state:   0x%x\n", status.state);
	printf("\n\n");

err:
	close(fd);
	return ret;
}

static int do_rtmr_start(void)
{
	int fd;
	int ret;
	struct csv_rtmr_req req = { 0, };
	struct csv_guest_user_rtmr_start start = { 0, };

	// access csv-guest
	fd = open_csv_guest();
	if (fd < 0)
		return -1;

	start.version = CSV_RTMR_VERSION_MIN;

	// issue RTMR_EXTEND request
	req.buf = (__u64)&start;
	req.len = (__u64)sizeof(start);
	req.subcmd_id = CSV_GUEST_USER_RTMR_START;
	ret = ioctl(fd, CSV_CMD_RTMR, &req);

	// check return code
	if (ret == -ENODEV || ret == -ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		goto err;
	}

	if (errno == ENODEV || errno == ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		ret = -errno;
		goto err;
	}

	if (req.fw_error_code) {
		PR_ERR("fw_err: 0x%x\n", req.fw_error_code);
		ret = -EIO;
	}

err:
	close(fd);
	return ret;
}

static int do_rtmr_read(void)
{
	int fd;
	int ret;
	int i, j, k;
	int cmd_len;
	struct csv_rtmr_req req = { 0, };
	struct csv_guest_user_rtmr_read *cmd = NULL;

	// access csv-guest
	fd = open_csv_guest();
	if (fd < 0)
		return -1;

	cmd_len = sizeof(*cmd) + (CSV_RTMR_REG_NUM - 1) * CSV_RTMR_REG_SIZE;
	cmd = calloc(1, cmd_len);
	if (!cmd) {
		PR_ERR("cannot allocate EXTEND_READ command\n");
		ret = -1;
		goto err;
	}

	cmd->bitmap = 0xFFFFFFFFU;

	// issue RTMR_EXTEND request
	req.buf = (__u64)cmd;
	req.len = (__u64)cmd_len;
	req.subcmd_id = CSV_GUEST_USER_RTMR_READ;
	ret = ioctl(fd, CSV_CMD_RTMR, &req);

	// check return code
	if (ret == -ENODEV || ret == -ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		goto err;
	}

	if (errno == ENODEV || errno == ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		ret = -errno;
		goto err;
	}

	if (req.fw_error_code) {
		PR_ERR("fw_err: 0x%x\n", req.fw_error_code);
		ret = -EIO;
	}

	j = 0;
	PR_INFO("RTMR_READ values:\n");
	for (i = 0; i < CSV_RTMR_REG_NUM; i++) {
		if ((cmd->bitmap & (1 << i)) == 0)
			continue;

		PR_INFO("  RTMR[%02d]: ", i);
		for (k = 0; k < CSV_RTMR_REG_SIZE; k++) {
			printf("%02x", cmd->data[j * CSV_RTMR_REG_SIZE + k]);
		}
		printf("\n");
		j++;
	}
	printf("\n\n");
err:
	if (cmd)
		free(cmd);
	close(fd);
	return ret;
}

static int do_rtmr_extend(const char *str, int index)
{
	int fd;
	int ret;
	int len = strlen(str);
	struct csv_rtmr_req req = { 0, };
	struct csv_guest_user_rtmr_extend cmd = { 0, };

	// access csv-guest
	fd = open_csv_guest();
	if (fd < 0)
		return -1;

	// Each char in the @str correspond to 4 bits. The end of @str is '\0'.
	if (len != (2 * CSV_RTMR_EXTEND_LEN)) {
		PR_ERR("invalid length of input data to be extended, the length must be %d != %d\n",
			2 * CSV_RTMR_EXTEND_LEN, len);
		ret = -1;
		goto err;
	}

	// This demo only allows index in the range [3, CSV_RTMR_REG_INDEX_MAX].
	if (index < 3 || index > CSV_RTMR_REG_INDEX_MAX) {
		PR_ERR("the index requested is invalid, the valid index should be in the range [3, %d]\n",
			CSV_RTMR_REG_INDEX_MAX);
		ret = -1;
		goto err;
	}

	cmd.index = index;
	cmd.data_len = CSV_RTMR_EXTEND_LEN;
	ret = convert_ascii_str_to_bin(cmd.data, cmd.data_len, str, len);
	if (ret) {
		PR_ERR("fail to convert input data to binary\n");
		goto err;
	}

	// issue RTMR_EXTEND request
	req.buf = (__u64)&cmd;
	req.len = (__u64)sizeof(cmd);
	req.subcmd_id = CSV_GUEST_USER_RTMR_EXTEND;
	ret = ioctl(fd, CSV_CMD_RTMR, &req);

	// check return code
	if (ret == -ENODEV || ret == -ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		goto err;
	}

	if (errno == ENODEV || errno == ENOTTY) {
		PR_ERR("RTMR unsupported !!!\n");
		ret = -errno;
		goto err;
	}

	if (req.fw_error_code) {
		PR_ERR("fw_err: 0x%x\n", req.fw_error_code);
		ret = -EIO;
	}

err:
	close(fd);
	return ret;
}

static int _read_ccel(uint64_t* log_start, uint64_t* log_end)
{
	FILE* f = fopen(CCEL_TABLE, "rb");

	CCEL ccel = { 0, };

	if (!f) {
		PR_ERR("failed to open %s\n", CCEL_TABLE);
		return -1;
	}

	fread(&ccel, sizeof(CCEL), 1, f);
	fclose(f);

	if (strncmp(ccel.Signature, CCEL_SIG, strlen(CCEL_SIG)) != 0) {
		PR_ERR("invalid CCEL signature\n");
		return -1;
	}

	*log_start = ccel.LogStart;
	*log_end = ccel.LogStart + ccel.LogLen;

	PR_INFO("CCEL: start=0x%lx, len=0x%lx\n", ccel.LogStart, ccel.LogLen);
	return 0;
}

static uint8_t* _read_event_log(uint64_t start, size_t size)
{
	int fd = open(DEV_MEM, O_RDONLY);
	off_t page_base = 0;
	size_t offset = 0;
	size_t map_size = 0;
	uint8_t *map = NULL;
	uint8_t *log = malloc(size);

	if (fd < 0) {
		PR_ERR("failed to open %s\n", DEV_MEM);
		return NULL;
	}

	page_base = start & ~(PAGE_SIZE - 1);
	offset = start - page_base;
	map_size = offset + size;
	map = mmap(NULL, map_size, PROT_READ, MAP_SHARED, fd, page_base);
	if (map == MAP_FAILED) {
		PR_ERR("failed to mmap 0x%lx\n", page_base);
		close(fd);
		return NULL;
	}

	log = malloc(size);
	if (!log) {
		PR_ERR("failed to allocate log buffer\n");
		goto err_unmap;
	}

	memcpy(log, map + offset, size);

err_unmap:
	munmap(map, map_size);
	close(fd);
	return log;
}

static void _replay_event_log(uint8_t *log, size_t size)
{
	int i, j;
	size_t offset = 0;

	uint8_t pcr[CSV_RTMR_REG_NUM][CSV_RTMR_REG_SIZE] = { 0, }; // SM3 PCR
	uint8_t new_pcr[CSV_RTMR_REG_SIZE] = { 0, };

	TCG_PCR_EVENT_HDR *evt_hdr = (TCG_PCR_EVENT_HDR *)log;
	TCG_PCR_EVENT2 *evt = NULL;
	EVP_MD_CTX *ctx = NULL;

	// Skip SpecID event (EV_NO_ACTION)
	if (evt_hdr->EventType == EV_NO_ACTION) {
		offset += sizeof(TCG_PCR_EVENT_HDR) + evt_hdr->EventSize;
	}

	while ((offset + sizeof(TCG_PCR_EVENT2)) < size) {
		evt = (TCG_PCR_EVENT2 *)(log + offset);

		if (offset + sizeof(TCG_PCR_EVENT2) + evt->EventSize > size)
			break;

		if (evt->digests.count < 1 ||
		    evt->digests.digests[0].hashAlg != RTMR_HASH_ALG_SM3_256) {
			PR_ERR("Unsupported hash algorithm or empty digest\n");
			break;
		}

		ctx = EVP_MD_CTX_new();
		EVP_DigestInit_ex(ctx, EVP_sm3(), NULL);
		EVP_DigestUpdate(ctx, pcr[evt->PCRIndex], CSV_RTMR_REG_SIZE);
		EVP_DigestUpdate(ctx, evt->digests.digests[0].digest, CSV_RTMR_REG_SIZE);
		EVP_DigestFinal_ex(ctx, new_pcr, NULL);
		EVP_MD_CTX_free(ctx);

		memcpy(pcr[evt->PCRIndex], new_pcr, CSV_RTMR_REG_SIZE);
		offset += sizeof(TCG_PCR_EVENT2) + evt->EventSize;
	}

	PR_INFO("Final RTMR (SM3) value: \n");
	for (i = 0; i < CSV_RTMR_REG_NUM; i++) {
		PR_INFO("  RTMR[%02d] ", i);
		for (j = 0; j < 32; j++)
			printf("%02x", pcr[i][j]);
		printf("\n");
	}
	printf("\n\n");
}

static int do_rtmr_replay(void)
{
	uint64_t log_start, log_end;
	size_t log_size;
	uint8_t *log = NULL;
	int ret = 0;

	ret = _read_ccel(&log_start, &log_end);
	if (ret)
		return ret;

	log_size = log_end - log_start;
	log = _read_event_log(log_start, log_size);
	if (!log) {
		PR_ERR("failed to read event log\n");
		return -1;
	}

	_replay_event_log(log, log_size);

	free(log);
	return 0;
}

static int do_rtmr_dump(void)
{
	uint64_t log_start, log_end;
	size_t log_size;
	uint8_t *log = NULL;
	FILE *out = NULL;
	int ret = 0;

	ret = _read_ccel(&log_start, &log_end);
	if (ret)
		return ret;

	log_size = log_end - log_start;
	log = _read_event_log(log_start, log_size);
	if (!log) {
		PR_ERR("failed to read event log\n");
		return -1;
	}

	out = fopen(EVENT_LOG_OUT_FILE, "wb");
	if (!out) {
		PR_ERR("failed to open output file %s\n", EVENT_LOG_OUT_FILE);
		ret = -1;
		goto err_free_log;
	}

	fwrite(log, 1, log_size, out);
	fclose(out);

	PR_INFO("Event log saved to %s\n", EVENT_LOG_OUT_FILE);

err_free_log:
	free(log);
	return ret;
}

static struct option long_options[] = {
	{ "help", no_argument, 0, 'h' },
	{ "status", no_argument, 0, 's' },
	{ "start", no_argument, 0, 'S' },
	{ "read", no_argument, 0, 'r' },
	{ "extend", required_argument, 0, 'e' },
	{ "index", required_argument, 0, 'i' },
	{ "dump", no_argument, 0, 'd' },
	{ "replay", no_argument, 0, 'R' },
	{ 0, 0, 0, 0 },
};

static void print_rtmr_usage(void)
{
	printf("%-20s %-20s [--help] [--status] [--start] [--read]\n"
	       "%-42s[--extend <data> --index <index>] [--dump]\n"
	       "%-42s[--replay]\n\n",
	       g_prog_name, APP_SUBCMD_RTMR, "", "");
	printf("  --help:   Display this help message.\n");
	printf("  --status: Get RTMR status information.\n");
	printf("  --start:  Start RTMR (This can only execute once in the guest).\n");
	printf("  --read:   Read RTMR registers.\n");
	printf("  --extend: Extend RTMR register.\n"
	       "            You need provide ascii string follows --extend.\n\n"
	       "            For example:\n"
	       "                %s %s \\\n"
	       "                --extend 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef \\\n"
	       "                --index 3\n",
	       g_prog_name, APP_SUBCMD_RTMR);
	printf("  --index:  The index of the RTMR register to be extended.\n"
	       "            You need provide decimal index value follows --index.\n");
	printf("  --dump:   Dump Pre-OS coco event log to file %s.\n",
	       EVENT_LOG_OUT_FILE);
	printf("  --replay: Replay Pre-OS coco event log and re-calculate RTMR.\n");
}

int handle_subcmd_rtmr(int argc, char *argv[])
{
	int opt;

	bool is_rtmr_extend = false;
	int rtmr_index = INT_MAX;
	char *extend_data_str = NULL;
	char *endptr = NULL;

	if (argc < 2) {
		print_rtmr_usage();
		return -1;
	}

	while ((opt = getopt_long(argc, argv, "hsSre:i:dR", long_options, NULL)) != -1) {
		switch (opt) {
		case 'h':
			print_rtmr_usage();
			return 0;
		case 's':
			return do_rtmr_status();
		case 'S':
			return do_rtmr_start();
		case 'r':
			return do_rtmr_read();
		case 'e':
			is_rtmr_extend = true;
			extend_data_str = optarg;
			break;
		case 'i':
			rtmr_index = strtol(optarg, &endptr, 10);
			if (*endptr != '\0') {
				PR_ERR("Invalid rtmr index: %s\n", optarg);
				return -1;
			}
			break;
		case 'd':
			return do_rtmr_dump();
		case 'R':
			return do_rtmr_replay();
		default:
			PR_ERR("Invalid options !!!\n");
			print_rtmr_usage();
			return -1;
		}
	}

	if (is_rtmr_extend) {
		return do_rtmr_extend(extend_data_str, rtmr_index);
	}

	return 0;
}
